#!/usr/bin/env python
import numpy as np
import sys
import os 
import matplotlib.pyplot as plt
import scienceplots
import argparse

# # generate uav position and orientation
# if os.name == 'nt':  # Windows
#     lib_path = r"c:\Users\ljn\Documents\papers\Dual UWB localization and analysis\simulation"
# elif os.name == 'posix':  # macOS 或 Linux
#     lib_path = "/Users/lvjunning/Nutstore Files/papers/dual UWB localization and analysis/simulation"
# else:
#     raise EnvironmentError("Unsupported Operating System")
# # print(f"Library path: {lib_path}")
# sys.path.append(lib_path)

from range_estimator_4d import get_crlb, get_measurement, get_rotation, heading_filter_tightly_couple
from my_common_lib.matplot_common import linestyle, marker, labelfont, titlefont, color
# from range_estimator_4d import points_v1, get_traj, get_sparse_indices


from range_estimator_4d import  get_traj_pose as get_traj
from range_estimator_4d import  get_sparse_pose_indices as get_sparse_indices

# from range_estimator_4d import points_v1_for_interpolation as points_v1
from range_estimator_4d import points_v3 as points_v1

points = None
path_points_len = int(2000) #（3.3+4.0+3.3）/0.1*50

current_dir = None 

anchors_num = None
anchors = None
tags = None
path_points = None
orient = None
measurements = []
measurements_noise = []
crlb_values = []
Zrv = []
V = []



# measurements_np = np.asarray(measurements).reshape(len(measurements), tags.shape[0]*anchors.shape[0])
# measurements_noise_np = np.asarray(measurements_noise).reshape(len(measurements_noise), tags.shape[0]*anchors.shape[0])
# plt.figure(figsize=(12, 8))
# for i in range(6):
#     plt.plot(measurements_np[:, i], label="$r_{}$".format(i), linestyle="--")
#     plt.plot(measurements_noise_np[:, i], label="$\\hat{{r}}_{{{}}}$".format(i), alpha=0.7)
# plt.ylabel("distance/m")
# plt.legend()
# plt.grid(True)
# plt.tight_layout()
# plt.show()

# plot the true and estimation trajectory, CRLB and Zrv in separate plots
def plot_result(estimation_result_np):
    global V, Zrv
    Zrv_np = np.asarray(Zrv)
    V_np = np.asarray(V)
    time_stick = np.arange(path_points_len) * 40 / path_points_len
    index_to_time_factor = 40 / path_points_len
    with plt.style.context(['science','ieee']):
        cm = 1/2.54 # centimeters in inches
        plt.figure(figsize=(8.5*cm,6*cm*3/4),dpi = 600)
        label = ["$x$", "$y$", "$z$", "$\gamma$"]
        hat_label = ["$\hat{x}$", "$\hat{y}$", "$\hat{z}$", "$\hat{\gamma}$"]
        for i in range(3):
            plt.plot(time_stick, estimation_result_np[:, i], label=hat_label[i], color = color[i], alpha=0.7, linestyle="--", linewidth=0.9)
        for i in range(3):
            plt.plot(time_stick, path_points[:, i], label=label[i], color = color[i], linestyle="-", linewidth=0.7)
        plt.ylabel("Position/m")
        
        if anchors_num == 2:
            plt.ylim([-1.0, 5.0])
            plt.fill_between(x=np.arange(len(V_np)) *index_to_time_factor, y1=-1, y2=5, where=np.abs(V_np) < 0.4, color='red', alpha=0.2, edgecolor='none')
        plt.legend(ncols=2)
        # plt.grid(True)
        plt.savefig(os.path.join(current_dir, "localizaiton_result_{}_anchors_position_plot.png".format(anchors_num)), dpi=600)
        plt.close()

        # plt.subplot(4, 1, 2)
        plt.figure(figsize=(8.5 * cm, 6 * cm*3/4), dpi=600)
        plt.plot(time_stick, estimation_result_np[:, 3], color = color[3], label=hat_label[3], alpha=0.7, linestyle="--", linewidth=0.9)
        plt.plot(time_stick, orient[:, 0], label=label[3], color = color[3], linestyle="-", linewidth=0.7)
        plt.ylabel("Yaw/rad")
        plt.ylim([-0.3, 0.3])
        plt.legend()
        # plt.grid(True)
        plt.savefig(os.path.join(current_dir, "localizaiton_result_{}_anchors_yaw_plot.png".format(anchors_num)), dpi=600)
        plt.close()

        plt.figure(figsize=(8.5 * cm, 6 * cm*3/4), dpi=600)
        crlb_values_np = np.asarray(crlb_values).reshape(len(crlb_values), 4)
        plt.plot(time_stick, crlb_values_np[:, 0], label="$c_{x}$", color = color[0], linestyle="--")
        plt.plot(time_stick, crlb_values_np[:, 1], label="$c_y$", color = color[1], linestyle="--")
        plt.plot(time_stick, crlb_values_np[:, 2], label="$c_z$", color = color[2], linestyle="--")
        plt.plot(time_stick, crlb_values_np[:, 3], label="$c_\gamma$", color = color[3], linestyle="--")
        sparse_indices = get_sparse_indices(points, number=path_points_len)
        for idx in sparse_indices:
            plt.axvline(x=idx * index_to_time_factor, color='gray', linestyle=':', linewidth=1)
        # plot the Zrv and V
        # plt.yscale('log')
        plt.legend(ncols=4)
        # plt.grid(True)
        plt.ylabel("CRLB Diagonal Element")
        if anchors_num != 2:
            plt.xlabel("Time (s)")
        if anchors_num == 2:
            # plt.ylim([0, 1e8])
            plt.fill_between(x=np.arange(len(V_np))  * index_to_time_factor  ,y1=0,  y2=1e8, where=np.abs(V_np) < 0.4, color='red', alpha=0.2, edgecolor='none')
            plt.xlabel("Time (s)")
        plt.savefig(os.path.join(current_dir, "localizaiton_result_{}_anchors_crlb_plot.png".format(anchors_num)), dpi=600)
        plt.close()

        # crlb two axes
        # plt.figure(figsize=(8.5 * cm, 6 * cm*3/4), dpi=600)
        fig, ax1 = plt.subplots(figsize=(8.5*cm,6*cm*3/4), dpi=600)
        ax2 = ax1.twinx()
        crlb_values_np = np.asarray(crlb_values).reshape(len(crlb_values), 4)
        ax1.plot(time_stick, crlb_values_np[:, 0], label="$c_{x}$", color = color[0], linestyle="--")
        ax1.plot(time_stick, crlb_values_np[:, 1], label="$c_y$", color = color[1], linestyle="--")
        ax1.plot(time_stick, crlb_values_np[:, 2], label="$c_z$", color = color[2], linestyle="--")
        ax2.plot(time_stick, crlb_values_np[:, 3], label="$c_\gamma$", color = color[3], linestyle="--")
        sparse_indices = get_sparse_indices(points, number=path_points_len)
        for idx in sparse_indices:
            plt.axvline(x=idx * index_to_time_factor, color='gray', linestyle=':', linewidth=1)
        handles1, labels1 = ax1.get_legend_handles_labels()
        handles2, labels2 = ax2.get_legend_handles_labels()
        ax1.legend(handles1, labels1, loc='upper right', ncols = 1, bbox_to_anchor=(0.2, 0.7), fontsize='small')
        ax2.legend(handles2, labels2, loc='upper right', bbox_to_anchor=(0.2, 0.4), fontsize='small')
        ax2.set_ylabel("Yaw CRLB Component ($\\text{rad}^2$)")
        ax1.set_ylabel("Position CRLB Component ($\\text{m}^2$)")
        ax1.set_xlabel("Time (s)")
        if anchors_num != 2:
            ax1.fill_between(x=np.arange(len(V_np))  * index_to_time_factor  ,y1=0,  y2=1e8, where=np.abs(V_np) < 0.4, color='red', alpha=0.2, edgecolor='none')
        plt.savefig(os.path.join(current_dir, "localizaiton_result_{}_anchors_crlb_plot_two_axes.png".format(anchors_num)), dpi=600)
        plt.close()

        if anchors_num == 2:
            plt.figure(figsize=(8.5 * cm, 6 * cm*3/4), dpi=600)
            # plt.subplot(4, 1, 4)
            plt.plot(time_stick, Zrv_np, label="$z_{\mathbf{r}_v}$", color = color[6], linestyle="-")
            plt.plot(time_stick, V_np, label="$V$", color = color[5], linestyle="-")
            plt.axhline(y=0, color='gray', linestyle='--', linewidth=0.5)
            # plt.ylim([-4, 4])
            
            # plt.fill_between(x=np.arange(len(V_np))  * index_to_time_factor, y1=-4,  y2=4, where=np.abs(V_np) < 0.4, color='red', alpha=0.2, label="$\lvert V \\rvert < 0.2$",edgecolor='none')
            plt.legend()
            # plt.grid(True)
            plt.ylabel("Element of Determinant")
            plt.xlabel("Time (s)")
            plt.savefig(os.path.join(current_dir, "localizaiton_result_{}_anchors_zrv_v_plot.png".format(anchors_num)), dpi=600)
            plt.close()
        # plt.tight_layout()
        # plt.show()
        # plt.savefig(os.path.join(current_dir, "localization_result_{}_anchors.png".format(anchors_num)), dpi = 600)
        
        # calculate the the error
        position_error = estimation_result_np[:,0:3] - path_points[:,0:3]
        orientation_error = estimation_result_np[:,3] - path_points[:,5]
        # plot the error
        fig, ax1 = plt.subplots(figsize=(8.5*cm,6*cm*3/4), dpi=600)
        ax2 = ax1.twinx()
        for i in range(3):
            ax1.plot(time_stick, position_error[:, i], label=label[i], color = color[i], alpha=0.7, linestyle="-", linewidth=0.7)
        ax2.plot(time_stick, orientation_error, label=label[3], color = color[3], linestyle="-", linewidth=0.7)
        # line_gamma2 = ax2.plot(partial_heading_rmse, color='b', marker = marker[2], label='$\gamma_p$')
        ax1.set_ylabel('Error of Position (m)', family='Times New Roman')
        ax1.set_xlabel('Time (s)', family='Times New Roman')
        ax2.set_ylabel('Error of Yaw (rad)', family='Times New Roman')
        handles1, labels1 = ax1.get_legend_handles_labels()
        handles2, labels2 = ax2.get_legend_handles_labels()
        # ax1.legend(handles1, labels1, loc='upper right', bbox_to_anchor=(0.6, 1), fontsize='small')
        # ax2.legend(handles2, labels2, loc='upper right', bbox_to_anchor=(0.9, 1), fontsize='small')
        ax1.legend(handles1, labels1, loc='upper right', ncols = 3, bbox_to_anchor=(0.75, 0.2), fontsize='small')
        ax2.legend(handles2, labels2, loc='upper right', bbox_to_anchor=(0.9, 0.2), fontsize='small')
        # ax1.set_ylim(0.05, 0.22)
        # ax2.set_ylim(0.03, 0.07)
        plt.tight_layout()
        plt.savefig(os.path.join(current_dir, "localizaiton_result_{}_anchors_error_plot.png".format(anchors_num)), dpi=600)
        plt.close()
        plt.show()
# plot the true and estimation trajectory, CRLB and Zrv in one plot
def plot_result_style_1(estimation_result_np):
    global V, Zrv
    Zrv_np = np.asarray(Zrv)
    V_np = np.asarray(V)
    with plt.style.context(['science','ieee']):
        cm = 1/2.54 # centimeters in inches
        plt.figure(figsize=(8.5*cm,6*cm*3),dpi = 600)
        plt.subplot(4, 1, 1)
        label = ["$x$", "$y$", "$z$", "$\gamma$"]
        hat_label = ["$\hat{x}$", "$\hat{y}$", "$\hat{z}$", "$\hat{\gamma}$"]
        for i in range(3):
            plt.plot(estimation_result_np[:, i], label=hat_label[i], color = color[i], alpha=0.7, linestyle="--", linewidth=0.9)
        for i in range(3):
            plt.plot(path_points[:, i], label=label[i], color = color[i], linestyle="-", linewidth=0.7)
        plt.ylabel("Position/m")
        
        if anchors_num == 2:
            plt.ylim([-1.0, 5.0])
            plt.fill_between(x=range(len(V_np)),y1=-1,  y2=5, where=np.abs(V_np) < 0.4, color='red', alpha=0.2, edgecolor='none')
        plt.legend(ncols=2)
        # plt.grid(True)

        plt.subplot(4, 1, 2)
        plt.plot(estimation_result_np[:, 3], color = color[3], label=hat_label[3], alpha=0.7, linestyle="--", linewidth=0.9)
        plt.plot(orient[:, 2], label=label[3], color = color[3], linestyle="-", linewidth=0.7)
        plt.ylabel("Yaw/rad")
        # plt.ylim([-0.3, 0.3])
        plt.legend()
        # plt.grid(True)

        plt.subplot(4, 1, 3)
        crlb_values_np = np.asarray(crlb_values).reshape(len(crlb_values), 4)
        plt.plot(crlb_values_np[:, 0], label="$\sigma^2_x$", color = color[0], linestyle="--")
        plt.plot(crlb_values_np[:, 1], label="$\sigma^2_y$", color = color[1], linestyle="--")
        plt.plot(crlb_values_np[:, 2], label="$\sigma^2_z$", color = color[2], linestyle="--")
        plt.plot(crlb_values_np[:, 3], label="$\sigma^2_\gamma$", color = color[3], linestyle="--")
        sparse_indices = get_sparse_indices(points, number=path_points_len)
        for idx in sparse_indices:
            plt.axvline(x=idx, color='gray', linestyle=':', linewidth=1)
        # plot the Zrv and V
        plt.yscale('log')
        plt.legend(ncols=4)
        # plt.grid(True)
        plt.ylabel("CRLB Diagonal Element")
        if anchors_num != 2:
            plt.xlabel("Index")
        if anchors_num == 2:
            plt.ylim([0, 1e8])
            plt.fill_between(x=np.arange(len(V_np)) ,y1=0,  y2=1e8, where=np.abs(V_np) < 0.4, color='red', alpha=0.2, edgecolor='none')

        if anchors_num == 2:
            plt.subplot(4, 1, 4)
            plt.plot(Zrv_np, label="$z_{\mathbf{r}_v}$", color = color[6], linestyle="-")
            plt.plot(V_np, label="$V$", color = color[5], linestyle="-")
            plt.axhline(y=0, color='gray', linestyle='-', linewidth=0.5)
            plt.ylim([-4, 4])
            
            plt.fill_between(x=range(len(V_np)),y1=-4,  y2=4, where=np.abs(V_np) < 0.4, color='red', alpha=0.2, label="$\lvert V \\rvert < 0.2$",edgecolor='none')
            plt.legend()
            # plt.grid(True)
            plt.ylabel("Element of Determinant")
            plt.xlabel("Index")

        # plt.tight_layout()
        # plt.show()
        plt.savefig(os.path.join(current_dir, "localization_result_{}_anchors.png".format(anchors_num)), dpi = 600)
def save_result(result):
    est_file_name = os.path.join(current_dir, "localization_result_{}_anchors.txt".format(anchors_num))
    i = 1
    while os.path.exists(est_file_name):
        est_file_name = os.path.join(current_dir, "localization_result_{}_anchors_{}.txt".format(anchors_num, i))
        i += 1
    np.savetxt(est_file_name, result)

    gt_file_name = os.path.join(current_dir, "localization_result_{}_anchors_gt.txt".format(anchors_num))
    if (os.path.exists(gt_file_name) == False):
        np.savetxt(gt_file_name, path_points)

def ekf_loop():
    estimation_result= []
    state_init = np.array([points[0,0], points[0,1], points[0,2], 0.0])
    print("state_init: {}".format(state_init))
    ekf_filter = heading_filter_tightly_couple(anchors, tags, state_init)
    for i in range(len(measurements_noise)):
        ekf_filter.predict()
        ekf_filter.roll = orient_noise[i,0]
        ekf_filter.pitch = orient_noise[i,1]
        ekf_filter.update(measurements_noise[i])
        # ekf_filter.roll = orient[i,0]
        # ekf_filter.pitch = orient[i,1]
        # ekf_filter.update(measurements[i])
        estimation_result.append(ekf_filter.state)
    # measurements_np = np.asarray(measurements).reshape(len(measurements), tags.shape[0]*anchors.shape[0])
    estimation_result_np = np.asarray(estimation_result).reshape(len(estimation_result), 4)
    plot_result(estimation_result_np)
    # plot_result_style_1(estimation_result_np)
    save_result(estimation_result_np)


if __name__ == '__main__':

    # 解析命令行参数
    parser = argparse.ArgumentParser(description="Start the Optimization.")
    parser.add_argument('--anchors_config_path', type=str, default='None', help='anchor config file(default: None)')
    args = parser.parse_args()
    

    anchors = np.loadtxt(args.anchors_config_path)
    anchors_num = anchors.shape[0]
    current_dir = os.path.dirname(args.anchors_config_path)
    print('current_dir: ', current_dir)

    tags = np.array([[0, 0.2, -0.27],
                        [0, -0.2, -0.27]])

    points = points_v1[2:5] # [2:5] # [2:5]
    path_points = get_traj(sparse_traj=points, number=path_points_len) #（3.3+4.0+3.3）/0.1*50

    # calculate the true range measurements
    for i in range(len(path_points)):
        measurements.append(get_measurement(path_points[i,0:3], path_points[i,3:6], anchors, tags))

    # add noise to the roll and pitch, and the range measurements
    orient = path_points[:,3:6]
    orient_noise = orient + np.random.normal(0, 0.01, orient.shape)
    for i in range(len(measurements)):
        measurements_noise.append(measurements[i] + np.random.normal(0, 0.1, measurements[i].shape))

    for path_point in path_points:
        crlb = get_crlb(anchors, tags, path_point)
        # if crlb[0, 0] > 1e5:
        #     crlb = np.zeros((4,4))
        crlb_values.append([crlb[0, 0], crlb[1, 1], crlb[2, 2], crlb[3, 3]])

    # calucate the Zrv and V
    if anchors_num ==2:
        for i in range(len(path_points)):
            Rot = get_rotation(orient[i,0], orient[i,1], orient[i,2])
            ra1_b = np.dot(Rot.T , anchors[0,:] - path_points[i,0:3]) 
            ra2_b = np.dot(Rot.T , anchors[1,:] - path_points[i,0:3]) 
            rn1 = tags[0,:]
            rn2 = tags[1,:]
            rv = np.cross(ra1_b - ra2_b, rn1 - rn2)
            V_i = np.dot(ra1_b - ra2_b, np.cross(ra2_b - rn2, rn1 - rn2))
            Zrv.append(rv[2])
            V.append(V_i)

    ekf_loop()
